import json
import time
from concurrent.futures import ThreadPoolExecutor

import numpy as np
from sql_options import get_data_trace_clients_data
import torch
import torch.nn as nn
import copy
import random
import hashlib
Hash = hashlib.sha256
import base64
import tenseal as ts
from pypbc import *
from utils.masked_input import Verify_agg_sign_2, decompose_trac_w_new, Verify_sign_1, Hhash, Flatten_paremeter, KeyGen, \
    PPGen,PPGen1


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(784, 10)  # 784，100， 10
        self.fc2 = nn.Linear(10, 10)

    def forward(self, x):
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = torch.relu(x)
        x = self.fc2(x)
        return x

class Trace:
    def __init__(self,args):
        self.args = args.copy
        # self.device = args.device
        self.device = args["device"]
        self.num_clients = 0
        self.seed = args["seed"]
        self.net_glob = Net().to(self.device)
        self.dataset_train = None
        self.context = None
        '''初始化'''
        self.context = ts.context(ts.SCHEME_TYPE.BFV, poly_modulus_degree=4096, plain_modulus=1032193)
        self.num_trace = 0

    def func_init(self, args, num_clients):
        random.seed(self.seed)
        np.random.seed(self.seed)
        torch.manual_seed(self.seed)
        torch.cuda.manual_seed_all(self.seed)
        self.net_glob.train()
        # copy weights
        self.w_glob = self.net_glob.state_dict()
        self.num_clients = num_clients
        '''初始化'''
        [self.G2, self.g2, self.params, self.pairing] = KeyGen(512, 160)
        with torch.no_grad():
            self.total_params = sum(p.numel() for p in self.net_glob.state_dict().values())
            print(self.total_params)

#通过多线程生成pp1
        # self.str_pp1 = PPGen1(self.g2, self.params, self.pairing, self.total_params)  # 参数长度 d
        # self.pp1 = [Element(self.pairing, G2, value=s) for s in self.str_pp1]
#直接生成pp1
        self.pp1 = PPGen(self.g2, self.params, self.pairing, self.total_params)  # 参数长度 d
        self.str_pp1 = [Element.__str__(elem) for elem in self.pp1]

        self.len_pp2 = self.num_clients * 130
        self.pp2 = PPGen(self.g2, self.params, self.pairing, self.len_pp2)
        self.pk_sign_2_s = [0 for i in range(self.num_clients)]
        self.str_g2 = Element.__str__(self.g2)
        self.str_pp2 = [Element.__str__(elem) for elem in self.pp2]

        print("init done")

    def get_client_params(self):
        client_context_bytes = self.context.serialize(save_public_key=True,
                                               save_secret_key=True,
                                               save_relin_keys=True,
                                               save_galois_keys=True)
        client_context_b64 = base64.b64encode(client_context_bytes).decode('utf-8')
        return {
            'context': client_context_b64,
            'g2': self.str_g2,
            'params': str(self.params),
            'total_params': self.total_params,
            'pp1': self.str_pp1,
            'pp2': self.str_pp2,
            'len_pp2': self.len_pp2
        }
    def get_server_params(self):
        agg_context_bytes = self.context.serialize(save_public_key=True,
                                                   save_secret_key=False,
                                                   save_relin_keys=True,
                                                   save_galois_keys=True)
        # 将序列化的数据转为可传输格式，比如 base64
        agg_context_b64 = base64.b64encode(agg_context_bytes).decode('utf-8')
        return {
            'context': agg_context_b64,
            'g2': self.str_g2,
            'params': str(self.params),
            'total_params': self.total_params,
        }

    def start_trace(self,connection):
        unhonest_participant = []
        pk_sign_1_s = []
        pk_sign_2_s = []
        pk_sign_1_s_trac, pk_sign_2_s_trac = [], []
        for pks in pk_sign_1_s:
            if pks != 0:
                pk_sign_1_s_trac.append(pks)

        tra_enc_w_new_s, trac_sig1_s, trac_sig2_s = [], [], []
        trac_decry_w_new_list = []
        num_user = self.num_clients

        all_clients = get_data_trace_clients_data(connection)
        for client in all_clients:
            pk_sign_1 = Element(self.pairing, self.G2, value=client['pk_sign_1'])
            pk_sign_1_s_trac.append(pk_sign_1)
            tra_enc_w_new_s.append(base64.b64decode(client['client_BFV']))
            trac_sig1_s.append(client['client_sign1'])
            client_sign2 = client['client_sign2']
            trac_sig2_idx = Element(self.pairing, self.G2, value=client_sign2)
            trac_sig2_s.append(trac_sig2_idx)
            pk_sign_1_s.append(pk_sign_1)
            pk_sign_2 = Element(self.pairing, self.G2, value=client['pk_sign_2'])
            pk_sign_2_s.append(pk_sign_2)

        verify_trace_2 = Verify_agg_sign_2(tra_enc_w_new_s, self.pairing, trac_sig1_s, pk_sign_2_s, num_user, self.g2,
                                           trac_sig2_s)
        print("verify_trace_2 = ", verify_trace_2)

        if verify_trace_2 == 1:
            # 解密
            for i in range(num_user):
                trac_enc_w_new_bfv = ts.bfv_vector_from(self.context, tra_enc_w_new_s[i])
                trac_decry_w_new = trac_enc_w_new_bfv.decrypt()
                trac_decry_w_new_list.append(trac_decry_w_new)

            decrypted_w_loc_list, h1_idx_trac_list, h1_idx_ASCII_list = [], [], []
            for idx_1, trac_decr_w_new in zip(range(self.num_clients), trac_decry_w_new_list):
                decrypted_w_loc, h1_idx_chr_str, h1_idx_ASCII = decompose_trac_w_new(trac_decr_w_new, idx_1,
                                                                                     self.len_pp2, self.num_clients)
                decrypted_w_loc_list.append(decrypted_w_loc)

                h1_idx_ASCII_list.append(h1_idx_ASCII)
                h1_idx_ele_trac = Element(self.pairing, self.G2, value=h1_idx_chr_str)
                h1_idx_trac_list.append(h1_idx_ele_trac)
            print("Trac_verify1 = ")
            len_hv = 130 * self.num_clients
            len_pp2 = len_hv

            for i in range(num_user):
                hv_ini = [0] * len_hv
                a = i * 130
                b = a + 130
                replace_indices = list(range(a, b))
                for k, j in enumerate(replace_indices):
                    hv_ini[j] = h1_idx_ASCII_list[i][k]

                h2_idx_pre = Hhash(self.pp2, hv_ini, len_pp2, self.pairing)
                h2_idx_trac = Element.__mul__(h1_idx_trac_list[i], h2_idx_pre)

                sig1_s_ele_trac = Element(self.pairing, self.G2, value=trac_sig1_s[i])
                print("Trac_verify1 = ")
                Trac_verify1 = Verify_sign_1(self.pairing, self.g2, sig1_s_ele_trac, pk_sign_1_s_trac[i], h2_idx_trac)
                print("Trac_verify1 = ",Trac_verify1)
                if Trac_verify1 == 0:
                    unhonest_participant.append(str(i))

            len_w_idx = len(decrypted_w_loc_list[0])
            if len(unhonest_participant) <= self.num_clients:
                for i in range(num_user):
                    h1_idx = Hhash(self.pp1, decrypted_w_loc_list[i], len_w_idx, self.pairing)
                    if h1_idx != h1_idx_trac_list[i]:
                        unhonest_participant.append(str(i))
                if len(unhonest_participant) == 0:
                    unhonest_participant.append("Server")

        else:
            unhonest_participant.append("Server")

        return unhonest_participant


